#include "winograd_4x4_to_36.h"

namespace mslite {

Winograd4x4To36::Winograd4x4To36(cl_context context, cl_device_id device, std::string clKernelName, std::vector<int> ishape): Operation(context) {
    this->clFileName = clKernelName + ".cl";
    this->clKernelName = clKernelName;
    this->context = context;
    this->inputNames = {"input"};
    this->inputShapes = {ishape};                                   // ishape must be nhwc
    this->outputNames = {"output"};
    std::vector<int> oshape = {ishape[0], 36, ishape[1] * ishape[2] / 16, ishape[3]};   // oshape is also nhwc
    this->outputShapes = {oshape};

    readSrcCode(this->clFileName, this->clCode);
    compileKernel(this->clCode, this->clCode.size(), this->clKernelName, this->context, device, this->program, this->kernel);
}

cl_int Winograd4x4To36::prepareMem(void) {
    cl_int status;
    std::vector<float> Bt = {
        1.0000000000f,  0.0000000000f, -2.5000004768f, -0.0000001192f,  1.0000001192f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f,  0.9428091049f,  1.3333333731f, -0.4714044929f, -0.6666667461f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f, -0.9428089857f,  1.3333334923f,  0.4714045525f, -0.6666667461f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f, -0.1178511307f, -0.0833333358f,  0.2357022613f,  0.1666666865f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f,  0.1178511307f, -0.0833333507f, -0.2357022911f,  0.1666666865f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f,  0.9999998808f, -0.0000000596f, -2.5000000000f,  0.0000000000f, 1.0000000000f, 0.0000000000f, 0.0000000000f, 
    };
    memObjs.clear();
    memObjs["bt"] = createImage2d(context, CL_FLOAT, 12, 1, Bt.data(), &status);
    int inputSize = getTotalSize(inputShapes[0]);
    int outputSize = getTotalSize(outputShapes[0]);
    memObjs["input"] = clCreateBuffer(context, CL_MEM_READ_ONLY |  CL_MEM_ALLOC_HOST_PTR, inputSize*sizeof(float), NULL, &status);
    memObjs["output"] = clCreateBuffer(context, CL_MEM_WRITE_ONLY | CL_MEM_ALLOC_HOST_PTR, outputSize*sizeof(float), NULL, &status);

    memTypes["bt"] = Image2d;
    memTypes["input"] = Buffer;
    memTypes["output"] = Buffer;

    return status;
}

cl_int Winograd4x4To36::setParams(void) {
    globalThreads.resize(3);
    localThreads.resize(3);
    localThreads = {1, 1, 1};
    globalThreads = {static_cast<size_t>(outputShapes[0][1] * outputShapes[0][2]), 6, static_cast<size_t>(outputShapes[0][3] / 4)};

    const std::vector<std::vector<size_t>> localCandidates = {{8, 6, 4}, {8, 6, 2}, {4, 6, 2},
                                                              {4, 6, 2}, {2, 6, 2}, {2, 6, 1},
                                                              {1, 6, 1}, {1, 3, 1}, {1, 1, 1}};
    int maxGrpSize = 512;
    for (auto candidate: localCandidates) {
        if (candidate[0]*candidate[1]*candidate[2] <= maxGrpSize) {
            localThreads = candidate;
            break;
        }
    }
    // std::cout << localThreads[0] << " " << localThreads[1] << " " << localThreads[2] << std::endl;
    return 0;
}

// cl_int Winograd4x4To36::getOutut(cl_command_queue commandQueue, std::vector<std::vector<float>> &output) {
//     cl_int status;
//     int osize = getTotalSize(this->outputShapes[0]);
//     output.clear();
//     output.resize(1);
//     output[0].resize(osize);
//     status = mslite::clEnqueueReadBuffer(commandQueue, memObjs["output"], CL_TRUE, 0, sizeof(float)*osize, output[0].data(), 0, NULL, NULL);
//     return status;
// }

Winograd4x4To36Naive::Winograd4x4To36Naive(cl_context context, cl_device_id device, std::vector<int> ishape): 
    Winograd4x4To36(context, device, std::string("winograd_4x4_to_36_naive"), ishape) {
}

cl_int Winograd4x4To36Naive::bind(void) {
    cl_int status;
    cl_int4 input_sliced_shape[] = {inputShapes[0][0], inputShapes[0][1], inputShapes[0][2], inputShapes[0][3]/4};
    cl_int4 output_sliced_shape[] = {outputShapes[0][0], outputShapes[0][1], outputShapes[0][2], outputShapes[0][3]/4};

    status = clSetKernelArg(kernel, 0, sizeof(cl_mem), static_cast<void *>(&memObjs["output"]));
    status = clSetKernelArg(kernel, 1, sizeof(cl_mem), static_cast<void *>(&memObjs["input"]));
    status = clSetKernelArg(kernel, 2, sizeof(cl_mem), static_cast<void *>(&memObjs["bt"]));
    status = clSetKernelArg(kernel, 3, sizeof(cl_int4), static_cast<void *>(input_sliced_shape));
    status = clSetKernelArg(kernel, 4, sizeof(cl_int4), static_cast<void *>(output_sliced_shape));
    return status;
}

Winograd4x4To36Tflite::Winograd4x4To36Tflite(cl_context context, cl_device_id device, std::vector<int> ishape):
    Winograd4x4To36(context, device, std::string("winograd_4x4_to_36_tflite"), ishape) {

}

cl_int Winograd4x4To36Tflite::bind(void) {
    cl_int status;
    cl_int4 a[] = {outputShapes[0][2], outputShapes[0][3]/4, inputShapes[0][2]/4, -1};
    cl_int4 b[] = {inputShapes[0][2], inputShapes[0][1], -1, outputShapes[0][1]};
    cl_int4 c[] = {outputShapes[0][2], 0, 0, 0};

    status = clSetKernelArg(kernel, 0, sizeof(cl_mem), static_cast<void *>(&memObjs["output"]));
    status = clSetKernelArg(kernel, 1, sizeof(cl_mem), static_cast<void *>(&memObjs["input"]));
    status = clSetKernelArg(kernel, 2, sizeof(cl_mem), static_cast<void *>(&memObjs["bt"]));
    status = clSetKernelArg(kernel, 3, sizeof(cl_int4), static_cast<void *>(a));
    status = clSetKernelArg(kernel, 4, sizeof(cl_int4), static_cast<void *>(b));
    status = clSetKernelArg(kernel, 5, sizeof(cl_int4), static_cast<void *>(c));
    return status;
}

Winograd4x4To36Image2d::Winograd4x4To36Image2d(cl_context context, cl_device_id device, std::vector<int> ishape): 
    Winograd4x4To36(context, device, std::string("winograd_4x4_to_36_image2d"), ishape) {
}

cl_int Winograd4x4To36Image2d::prepareMem(void) {
    cl_int status;
    std::vector<float> Bt = {
        1.0000000000f,  0.0000000000f, -2.5000004768f, -0.0000001192f,  1.0000001192f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f,  0.9428091049f,  1.3333333731f, -0.4714044929f, -0.6666667461f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f, -0.9428089857f,  1.3333334923f,  0.4714045525f, -0.6666667461f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f, -0.1178511307f, -0.0833333358f,  0.2357022613f,  0.1666666865f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f,  0.1178511307f, -0.0833333507f, -0.2357022911f,  0.1666666865f, 0.0000000000f, 0.0000000000f, 0.0000000000f, 
        0.0000000000f,  0.9999998808f, -0.0000000596f, -2.5000000000f,  0.0000000000f, 1.0000000000f, 0.0000000000f, 0.0000000000f, 
    };
    memObjs.clear();
    memObjs["bt"] = createImage2d(context, CL_FLOAT, 12, 1, Bt.data(), &status);
    int ih = inputShapes[0][1] * inputShapes[0][3] / 4, iw = inputShapes[0][2];
    int oh = outputShapes[0][1] * outputShapes[0][3] / 4, ow = outputShapes[0][2];
    memObjs["input"] = createImage2d(context, CL_FLOAT, iw, ih, NULL, &status);
    memObjs["output"] = createImage2d(context, CL_FLOAT, ow, oh, NULL, &status);

    memTypes["bt"] = Image2d;
    memTypes["input"] = Image2d;
    memTypes["output"] = Image2d;

    return status;
}

cl_int Winograd4x4To36Image2d::bind(void) {
    cl_int status;
    cl_int4 input_sliced_shape[] = {inputShapes[0][0], inputShapes[0][1], inputShapes[0][2], inputShapes[0][3]/4};
    cl_int4 output_sliced_shape[] = {outputShapes[0][0], outputShapes[0][1], outputShapes[0][2], outputShapes[0][3]/4};

    status = clSetKernelArg(kernel, 0, sizeof(cl_mem), static_cast<void *>(&memObjs["output"]));
    status = clSetKernelArg(kernel, 1, sizeof(cl_mem), static_cast<void *>(&memObjs["input"]));
    status = clSetKernelArg(kernel, 2, sizeof(cl_mem), static_cast<void *>(&memObjs["bt"]));
    status = clSetKernelArg(kernel, 3, sizeof(cl_int4), static_cast<void *>(input_sliced_shape));
    status = clSetKernelArg(kernel, 4, sizeof(cl_int4), static_cast<void *>(output_sliced_shape));
    return status;
}

}